//! Type system AST nodes for NRC language

use crate::error::Location;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// All possible types in NRC
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Type {
    /// Basic types
    Basic(BasicType),
    /// Generic type parameter
    Generic(String),
    /// Array/vector type
    Array(ArrayType),
    /// Map type
    Map(MapType),
    /// Tuple type
    Tuple(TupleType),
    /// Struct type
    Struct(StructType),
    /// Enum type
    Enum(EnumType),
    /// Function type
    Function(FunctionType),
    /// Reference type (&T)
    Reference(ReferenceType),
    /// Vec type (dynamic array)
    Vec(VecType),
    /// Channel type
    Channel(ChannelType),
    /// Optional type (nullable)
    Optional(OptionalType),
    /// Error type (errable)
    Error(ErrorType),
    /// Result type (built-in Result<T, E>)
    Result(ResultType),
    /// Option type (built-in Option<T>)
    Option(OptionType),
    /// Union type
    Union(UnionType),
    /// Type alias
    Alias(AliasType),
    /// Reference counted type (non-atomic)
    Rc(RcType),
    /// Weak reference type
    Weak(WeakType),
}

/// Basic types
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum BasicType {
    /// Integer types (Go-style naming)
    /// Default signed integer (64-bit)
    Int,
    /// 8-bit signed integer
    Int8,
    /// 16-bit signed integer
    Int16,
    /// 32-bit signed integer
    Int32,
    /// 64-bit signed integer
    Int64,
    /// Default unsigned integer (64-bit)
    Uint,
    /// 8-bit unsigned integer
    Uint8,
    /// 16-bit unsigned integer
    Uint16,
    /// 32-bit unsigned integer
    Uint32,
    /// 64-bit unsigned integer
    Uint64,

    /// Floating point types (Go-style naming)
    /// 32-bit floating point
    Float32,
    /// 64-bit floating point
    Float64,

    /// Boolean type
    /// Boolean value (true/false)
    Bool,

    /// String type
    /// String of characters
    String,

    /// Character type
    /// Single character (byte)
    Char,

    /// Rune type (Unicode code point)
    /// Unicode code point (int32 alias, like Go)
    Rune,

    /// Any type
    /// Any value type
    Any,

    /// Void type
    /// No value
    Void,
}

/// Array/vector type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ArrayType {
    /// Element type
    pub element_type: Box<Type>,
    /// Array size (if fixed size)
    pub size: Option<usize>,
    /// Location in source
    pub location: Location,
}

/// Map type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct MapType {
    /// Key type
    pub key_type: Box<Type>,
    /// Value type
    pub value_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Tuple type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TupleType {
    /// Element types
    pub element_types: Vec<Type>,
    /// Location in source
    pub location: Location,
}

/// Struct type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct StructType {
    /// Struct name
    pub name: String,
    /// Type arguments
    pub type_args: Vec<Type>,
    /// Location in source
    pub location: Location,
}

/// Enum type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct EnumType {
    /// Enum name
    pub name: String,
    /// Type arguments
    pub type_args: Vec<Type>,
    /// Location in source
    pub location: Location,
}

/// Function type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct FunctionType {
    /// Parameter types
    pub parameter_types: Vec<Type>,
    /// Return type
    pub return_type: Option<Box<Type>>,
    /// Is variadic
    pub variadic: bool,
    /// Location in source
    pub location: Location,
}

/// Reference type (&T) - mutable reference
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ReferenceType {
    /// Referenced type
    pub referenced_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Vec type (dynamic array)
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct VecType {
    /// Element type
    pub element_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Channel type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ChannelType {
    /// Element type
    pub element_type: Box<Type>,
    /// Channel direction
    pub direction: ChannelDirection,
    /// Buffer size (if buffered)
    pub buffer_size: Option<usize>,
    /// Location in source
    pub location: Location,
}

/// Channel direction
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ChannelDirection {
    /// Send only
    Send,
    /// Receive only
    Receive,
    /// Bidirectional
    Bidirectional,
}

/// Optional type (nullable)
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct OptionalType {
    /// Inner type
    pub inner_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Error type (errable)
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ErrorType {
    /// Inner type
    pub inner_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Result type (built-in Result<T, E>)
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ResultType {
    /// Ok type (success value type)
    pub ok_type: Box<Type>,
    /// Err type (error value type)
    pub err_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Option type (built-in Option<T>)
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct OptionType {
    /// Some type (value type)
    pub some_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Union type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct UnionType {
    /// Union member types
    pub member_types: Vec<Type>,
    /// Location in source
    pub location: Location,
}

/// Type alias
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AliasType {
    /// Alias name
    pub name: String,
    /// Type arguments
    pub type_args: Vec<Type>,
    /// Location in source
    pub location: Location,
}

/// Function parameter
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Parameter {
    /// Parameter name
    pub name: String,
    /// Parameter type
    pub param_type: Type,
    /// Default value (if any)
    pub default_value: Option<crate::ast::expr::Expression>,
    /// Location in source
    pub location: Location,
}

impl Type {
    /// Get the location of this type
    pub fn location(&self) -> Location {
        match self {
            Type::Basic(_) => Location::new(0, 0, 0),
            Type::Generic(_) => Location::new(0, 0, 0),
            Type::Array(ty) => ty.location,
            Type::Map(ty) => ty.location,
            Type::Tuple(ty) => ty.location,
            Type::Struct(ty) => ty.location,
            Type::Enum(ty) => ty.location,
            Type::Function(ty) => ty.location,
            Type::Reference(ty) => ty.location,
            Type::Vec(ty) => ty.location,
            Type::Channel(ty) => ty.location,
            Type::Optional(ty) => ty.location,
            Type::Error(ty) => ty.location,
            Type::Result(ty) => ty.location,
            Type::Option(ty) => ty.location,
            Type::Union(ty) => ty.location,
            Type::Alias(ty) => ty.location,
            Type::Rc(ty) => ty.location,
            Type::Weak(ty) => ty.location,
        }
    }

    /// Get the default zero value for this type
    /// Used when creating pointers with &Type{} without initializers
    pub fn default_value(&self) -> Option<crate::ast::expr::Literal> {
        use crate::ast::expr::Literal;
        match self {
            Type::Basic(basic) => Some(match basic {
                BasicType::Int
                | BasicType::Int8
                | BasicType::Int16
                | BasicType::Int32
                | BasicType::Int64 => Literal::Integer(0),
                BasicType::Uint
                | BasicType::Uint8
                | BasicType::Uint16
                | BasicType::Uint32
                | BasicType::Uint64 => Literal::Integer(0),
                BasicType::Float32 | BasicType::Float64 => Literal::Float(0.0),
                BasicType::Bool => Literal::Boolean(false),
                BasicType::String => Literal::String(String::new()),
                BasicType::Char => Literal::Char('\0'),
                BasicType::Rune => Literal::Integer(0),
                BasicType::Void | BasicType::Any => Literal::Null,
            }),
            Type::Reference(_) => Some(Literal::Null),
            Type::Optional(_) => Some(Literal::Null),
            Type::Result(_) => None, // Result types don't have default values
            Type::Option(_) => None, // Option types don't have default values
            _ => None,
        }
    }

    /// Check if this type is a basic type
    pub fn is_basic(&self) -> bool {
        matches!(self, Type::Basic(_))
    }

    /// Check if this type is a generic type parameter
    pub fn is_generic(&self) -> bool {
        matches!(self, Type::Generic(_))
    }

    /// Check if this type is a reference type
    pub fn is_reference(&self) -> bool {
        matches!(self, Type::Reference(_))
    }

    /// Check if this type is an array type
    pub fn is_array(&self) -> bool {
        matches!(self, Type::Array(_))
    }

    /// Check if the type (recursively) contains any generic parameters.
    pub fn contains_generic(&self) -> bool {
        match self {
            Type::Generic(_) => true,
            Type::Array(array) => array.element_type.contains_generic(),
            Type::Map(map) => map.key_type.contains_generic() || map.value_type.contains_generic(),
            Type::Tuple(tuple) => tuple.element_types.iter().any(Type::contains_generic),
            Type::Struct(struct_type) => struct_type.type_args.iter().any(Type::contains_generic),
            Type::Enum(enum_type) => enum_type.type_args.iter().any(Type::contains_generic),
            Type::Function(function_type) => {
                function_type
                    .parameter_types
                    .iter()
                    .any(Type::contains_generic)
                    || function_type
                        .return_type
                        .as_deref()
                        .map(Type::contains_generic)
                        .unwrap_or(false)
            }
            Type::Reference(reference_type) => reference_type.referenced_type.contains_generic(),
            Type::Vec(vec_type) => vec_type.element_type.contains_generic(),
            Type::Channel(channel_type) => channel_type.element_type.contains_generic(),
            Type::Optional(optional_type) => optional_type.inner_type.contains_generic(),
            Type::Error(error_type) => error_type.inner_type.contains_generic(),
            Type::Result(result_type) => {
                result_type.ok_type.contains_generic() || result_type.err_type.contains_generic()
            }
            Type::Option(option_type) => option_type.some_type.contains_generic(),
            Type::Union(union_type) => union_type.member_types.iter().any(Type::contains_generic),
            Type::Alias(alias_type) => alias_type.type_args.iter().any(Type::contains_generic),
            Type::Rc(rc_type) => rc_type.inner_type.contains_generic(),
            Type::Weak(weak_type) => weak_type.inner_type.contains_generic(),
            Type::Basic(_) => false,
        }
    }

    /// Check if this type is a function type
    pub fn is_function(&self) -> bool {
        matches!(self, Type::Function(_))
    }

    /// Check if this type is an optional type
    pub fn is_optional(&self) -> bool {
        matches!(self, Type::Optional(_))
    }

    /// Check if this type is an error type
    pub fn is_error(&self) -> bool {
        matches!(self, Type::Error(_))
    }

    /// Get the string representation of this type
    pub fn to_string(&self) -> String {
        match self {
            Type::Basic(basic) => match basic {
                BasicType::Int => "int".to_string(),
                BasicType::Int8 => "int8".to_string(),
                BasicType::Int16 => "int16".to_string(),
                BasicType::Int32 => "int32".to_string(),
                BasicType::Int64 => "int64".to_string(),
                BasicType::Uint => "uint".to_string(),
                BasicType::Uint8 => "uint8".to_string(),
                BasicType::Uint16 => "uint16".to_string(),
                BasicType::Uint32 => "uint32".to_string(),
                BasicType::Uint64 => "uint64".to_string(),
                BasicType::Float32 => "float32".to_string(),
                BasicType::Float64 => "float64".to_string(),
                BasicType::Bool => "bool".to_string(),
                BasicType::String => "string".to_string(),
                BasicType::Char => "char".to_string(),
                BasicType::Rune => "rune".to_string(),
                BasicType::Any => "any".to_string(),
                BasicType::Void => "void".to_string(),
            },
            Type::Generic(name) => name.clone(),
            Type::Array(ty) => {
                if let Some(size) = ty.size {
                    format!("[{}, {}]", ty.element_type.to_string(), size)
                } else {
                    format!("[{}]", ty.element_type.to_string())
                }
            }
            Type::Map(ty) => {
                format!(
                    "map<{}, {}>",
                    ty.key_type.to_string(),
                    ty.value_type.to_string()
                )
            }
            Type::Tuple(ty) => {
                let types: Vec<String> = ty.element_types.iter().map(|t| t.to_string()).collect();
                format!("({})", types.join(", "))
            }
            Type::Struct(ty) => {
                if ty.type_args.is_empty() {
                    ty.name.clone()
                } else {
                    let args: Vec<String> = ty.type_args.iter().map(|t| t.to_string()).collect();
                    format!("{}<{}>", ty.name, args.join(", "))
                }
            }
            Type::Enum(ty) => {
                if ty.type_args.is_empty() {
                    ty.name.clone()
                } else {
                    let args: Vec<String> = ty.type_args.iter().map(|t| t.to_string()).collect();
                    format!("{}<{}>", ty.name, args.join(", "))
                }
            }
            Type::Function(ty) => {
                let params: Vec<String> =
                    ty.parameter_types.iter().map(|t| t.to_string()).collect();
                let return_type = if let Some(rt) = &ty.return_type {
                    rt.to_string()
                } else {
                    "void".to_string()
                };
                format!("fn({}) -> {}", params.join(", "), return_type)
            }
            Type::Reference(ty) => {
                format!("&{}", ty.referenced_type.to_string())
            }
            Type::Vec(ty) => {
                format!("Vector<{}>", ty.element_type.to_string())
            }
            Type::Channel(ty) => {
                let direction = match ty.direction {
                    ChannelDirection::Send => "chan<-",
                    ChannelDirection::Receive => "<-chan",
                    ChannelDirection::Bidirectional => "chan",
                };
                if let Some(size) = ty.buffer_size {
                    format!("{} {}[{}]", direction, ty.element_type.to_string(), size)
                } else {
                    format!("{} {}", direction, ty.element_type.to_string())
                }
            }
            Type::Optional(ty) => {
                format!("{}?", ty.inner_type.to_string())
            }
            Type::Error(ty) => {
                format!("{}!", ty.inner_type.to_string())
            }
            Type::Result(ty) => {
                format!(
                    "Result<{}, {}>",
                    ty.ok_type.to_string(),
                    ty.err_type.to_string()
                )
            }
            Type::Option(ty) => {
                format!("Option<{}>", ty.some_type.to_string())
            }
            Type::Union(ty) => {
                let types: Vec<String> = ty.member_types.iter().map(|t| t.to_string()).collect();
                format!("({})", types.join(" | "))
            }
            Type::Alias(ty) => {
                if ty.type_args.is_empty() {
                    ty.name.clone()
                } else {
                    let args: Vec<String> = ty.type_args.iter().map(|t| t.to_string()).collect();
                    format!("{}<{}>", ty.name, args.join(", "))
                }
            }
            Type::Rc(ty) => {
                format!("rc<{}>", ty.inner_type.to_string())
            }
            Type::Weak(ty) => {
                format!("weak<{}>", ty.inner_type.to_string())
            }
        }
    }
}

/// Type context for type checking
#[derive(Debug, Clone)]
pub struct TypeContext {
    /// Type bindings
    pub bindings: HashMap<String, Type>,
    /// Generic type parameters
    pub generics: Vec<String>,
}

impl TypeContext {
    /// Create a new type context
    pub fn new() -> Self {
        Self {
            bindings: HashMap::new(),
            generics: Vec::new(),
        }
    }

    /// Add a type binding
    pub fn bind(&mut self, name: String, ty: Type) {
        self.bindings.insert(name, ty);
    }

    /// Get a type binding
    pub fn get(&self, name: &str) -> Option<&Type> {
        self.bindings.get(name)
    }

    /// Add a generic type parameter
    pub fn add_generic(&mut self, name: String) {
        self.generics.push(name);
    }

    /// Check if a name is a generic type parameter
    pub fn is_generic(&self, name: &str) -> bool {
        self.generics.contains(&name.to_string())
    }
}

impl Default for TypeContext {
    fn default() -> Self {
        Self::new()
    }
}

/// Reference counted type (non-atomic)
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct RcType {
    /// Inner type being reference counted
    pub inner_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

/// Weak reference type
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct WeakType {
    /// Inner type being weakly referenced
    pub inner_type: Box<Type>,
    /// Location in source
    pub location: Location,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_basic_types() {
        let int_type = Type::Basic(BasicType::Int);
        let string_type = Type::Basic(BasicType::String);
        let bool_type = Type::Basic(BasicType::Bool);

        assert_eq!(int_type.to_string(), "int");
        assert_eq!(string_type.to_string(), "string");
        assert_eq!(bool_type.to_string(), "bool");
    }

    #[test]
    fn test_array_type() {
        let element_type = Type::Basic(BasicType::Int);
        let location = Location::new(1, 1, 0);
        let array_type = Type::Array(ArrayType {
            element_type: Box::new(element_type),
            size: Some(10),
            location,
        });

        assert_eq!(array_type.to_string(), "[int, 10]");
        assert!(array_type.is_array());
    }

    #[test]
    fn test_reference_type() {
        let referenced_type = Type::Basic(BasicType::Int);
        let location = Location::new(1, 1, 0);
        let reference_type = Type::Reference(ReferenceType {
            referenced_type: Box::new(referenced_type),
            location,
        });

        assert_eq!(reference_type.to_string(), "&int");
        assert!(reference_type.is_reference());
    }

    #[test]
    fn test_function_type() {
        let param_types = vec![Type::Basic(BasicType::Int), Type::Basic(BasicType::String)];
        let return_type = Type::Basic(BasicType::Bool);
        let location = Location::new(1, 1, 0);
        let function_type = Type::Function(FunctionType {
            parameter_types: param_types,
            return_type: Some(Box::new(return_type)),
            variadic: false,
            location,
        });

        assert_eq!(function_type.to_string(), "fn(int, string) -> bool");
        assert!(function_type.is_function());
    }

    #[test]
    fn test_optional_type() {
        let inner_type = Type::Basic(BasicType::Int);
        let location = Location::new(1, 1, 0);
        let optional_type = Type::Optional(OptionalType {
            inner_type: Box::new(inner_type),
            location,
        });

        assert_eq!(optional_type.to_string(), "int?");
        assert!(optional_type.is_optional());
    }

    #[test]
    fn test_type_context() {
        let mut context = TypeContext::new();
        context.bind("x".to_string(), Type::Basic(BasicType::Int));
        context.add_generic("T".to_string());

        assert_eq!(context.get("x"), Some(&Type::Basic(BasicType::Int)));
        assert!(context.is_generic("T"));
        assert!(!context.is_generic("U"));
    }
}
